import numpy as np
from collections import Counter
import matplotlib.pyplot as plt

class KNN:
    def __init__(self, k=3):

        self.k = k
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):

        self.X_train = np.array(X)
        self.y_train = np.array(y)

    def euclidean_distance(self, x1, x2):

        return np.sqrt(np.sum((x1 - x2) ** 2))

    def predict(self, X):

        X = np.array(X)
        predictions = []

        for test_sample in X:
            # 计算当前测试样本与所有训练样本的距离
            distances = []
            for i, train_sample in enumerate(self.X_train):
                dist = self.euclidean_distance(test_sample, train_sample)
                distances.append((dist, self.y_train[i]))

            # 按距离排序并选择前k个最近邻
            distances.sort(key=lambda x: x[0])
            k_nearest = distances[:self.k]

            # 获取k个最近邻的标签
            k_nearest_labels = [label for _, label in k_nearest]

            # 投票决定预测标签（选择出现次数最多的标签）
            most_common = Counter(k_nearest_labels).most_common(1)
            predictions.append(most_common[0][0])

        return np.array(predictions)

    def predict_proba(self, X):

        X = np.array(X)
        probabilities = []

        for test_sample in X:
            distances = []
            for i, train_sample in enumerate(self.X_train):
                dist = self.euclidean_distance(test_sample, train_sample)
                distances.append((dist, self.y_train[i]))

            distances.sort(key=lambda x: x[0])
            k_nearest = distances[:self.k]
            k_nearest_labels = [label for _, label in k_nearest]

            # 计算每个类别的概率
            label_counts = Counter(k_nearest_labels)
            prob_dict = {}
            total = len(k_nearest_labels)

            for label in set(self.y_train):
                count = label_counts.get(label, 0)
                prob_dict[label] = count / total

            probabilities.append(prob_dict)

        return probabilities

    def accuracy(self, X_test, y_test):

        predictions = self.predict(X_test)
        return np.mean(predictions == y_test)